import numpy as np
import os

def gen_complex_gan_data_numpy():
    """
    生成 complex64 数据，但存储为两个 float 连续存储：
      - 输入: x.bin, y.bin (float32)
      - 输出: z.bin (float32), 实部虚部交替存储
        z = x + 1j*y, 存储格式: [x0, y0, x1, y1, ...]
    """
    np.random.seed(42)
    os.makedirs("./input", exist_ok=True)
    os.makedirs("./output", exist_ok=True)

    batch_size = 13
    features = 1525
    shape = (batch_size, features)

    # 1. 随机 float 输入
    x = (np.random.rand(*shape).astype(np.float32) * 2 - 1)  # [-1,1]
    y = (np.random.rand(*shape).astype(np.float32) * 2 - 1)

    # 2. 生成复数存储为 float32 两个连续元素
    z = np.empty((batch_size, features*2), dtype=np.float32)
    z[:, 0::2] = x
    z[:, 1::2] = y

    # 3. 保存文件
    x.tofile("./input/input_x.bin")
    y.tofile("./input/input_y.bin")
    z.tofile("./output/golden_z.bin")

    # 4. 抽查
    print("=== 抽查前4个 ===")
    print("x:", x.flatten()[:4])
    print("y:", y.flatten()[:4])
    # 4. 抽查：组合为复数输出 a+bi
    z_complex = z[:, 0::2] + 1j * z[:, 1::2]
    print("=== 抽查前4个复数 ===")
    for i in range(4):
        print(f"z[{i}] = {z_complex.flatten()[i]:.6f}")  # 自动 a+bj 格式

    print("\nshape:", z.shape, "dtype:", z.dtype)



if __name__ == "__main__":
    gen_complex_gan_data_numpy()
